# This is a sample Python script.

# Press Shift+F10 to execute it or replace it with your code.
# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings.

# 导入 efinance 如果没有安装则需要通过执行命令: pip install efinance 来安装
import efinance as ef

import json
import requests
from loguru import logger
import time


def _get_market_code(stock_code):
    """
    根据股票代码计算出市场代码。
    :param stock_code: 股票代码
    :return: 市场代码，0：深圳 北京，1：上海，2：其他
    """
    # 获取股票代码的前缀
    code_prefix = int(stock_code[:1])

    # 根据前缀判断市场
    if code_prefix in [0, 2, 3, 4, 8]:  # 深圳股票代码前缀一般为 0、2、3
        return 0  # 深圳市场
    elif code_prefix in [6, 9]:  # 上海股票代码前缀一般为 6、9
        return 1  # 上海市场
    else:
        return 2  # 其他市场，此处假设为北京市场


def _handle_event(event_data) -> str:
    """
    解析event数据
    :param event_data: SSE数据
    :return: 事件的数据字段
    """
    lines = event_data.strip().split('\n')
    data = ""

    for line in lines:
        if line.startswith("event:"):  # 东财分时数据的事件无event
            event_type = line.replace("event:", "").strip()
        elif line.startswith("data:"):
            data = line.replace("data:", "").strip()
    # logger.info(f"data received: {data}")
    return data


def get_minutely_data(code: str) -> str:
    """
    获取最新的分时数据
    :param code: 股票代码，6位数字格式，比如：“000001”。
    :return: 数据字典
    """
    # url格式，需要调用者依次填充：市场代码（0：深圳，0：上海），股票代码，日期窗口
    url_format = "http://16.push2.eastmoney.com/api/qt/stock/details/sse?fields1=f1,f2,f3,f4&fields2=f51,f52,f53,f54,f55&mpi=2000&ut=bd1d9ddb04089700cf9c27f6f7426281&fltt=2&pos=-0&secid={market}.{code}" \
                 "&ndays={days}&iscr=0&iscca=0&wbp2u=1849325530509956|0|1|0|web"
    res_format = "{market}|{code}|{jhjj}|{jhjj}"
    market = _get_market_code(code)
    url = url_format.format(market=market, code=code, days=1)  # 请求url
    headers = {
        "Accept": "text/event-stream"
    }
    print(url)
    response = requests.get(url, headers=headers, stream=True)  # 请求数据，开启流式传输

    if response.status_code == 200:
        event_data = ""
        for chunk in response.iter_content(chunk_size=None):  # 不断地获取数据
            data_decoded = chunk.decode('utf-8')  # 将字节流解码成字符串
            event_data += data_decoded
            parts = event_data.split('\n\n')
            if len(parts) > 1:  # 获取到一条完整的事件后，对数据进行解析。
                data = _handle_event(parts[0])
                l = json.loads(data).get("data").get("details")
                for v in l:
                    t = v[:7]
                    if t.startswith("09:25:0"):
                        return res_format.format(market=market, code=code, jhjj=v.split(",")[2])
                    elif int(v[:1]) == 1 or int(v[3:4]) > 2:
                        return res_format.format(market=market, code=code, jhjj=0)
            return ""


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    get_minutely_data("688720")
    codes = ef.stock.get_realtime_quotes()
    with open('output.txt', 'w') as f:
        # codes["股票代码"].apply(lambda x: print(get_minutely_data(x)))
        for index, code in codes.iterrows():
            print(code["股票代码"])
            jhjj = get_minutely_data(code["股票代码"])
            logger.info(f"data_test: {jhjj}")
            f.write(str(jhjj) + '\n')
            # time.sleep(1)
    logger.info("all done")
    # 股票代码
    stock_code = '600519'
    # 数据间隔时间为 1 分钟
    freq = 1
    # 获取最新一个交易日的分钟级别股票行情数据
    # df = ef.stock.get_quote_history(stock_code, klt=freq)
    # 将数据存储到 csv 文件中
    # df.to_csv(f'{stock_code}.csv', encoding='utf-8-sig', index=None)
    # print(f'股票: {stock_code} 的行情数据已存储到文件: {stock_code}.csv 中！')

# See PyCharm help at https://www.jetbrains.com/help/pycharm/
